package com.idatax.auth.common.jwt;

import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.ECKey;
import com.nimbusds.jose.jwk.RSAKey;
import io.jsonwebtoken.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.security.Key;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;

@Component
public class JwtUtils {

    @Value("${jwt.expiration}")
    private Long expiration;

    @Value("${jwt.issuer}")
    private String issuer;

    @Autowired
    private JwkKeyService jwkKeyService;

    public String generateToken(String username, Map<String, Object> customClaims) {
        JWK jwk = jwkKeyService.getSigningKey();
        PrivateKey privateKey;
        try {
            privateKey = jwk instanceof RSAKey ? 
                ((RSAKey)jwk).toPrivateKey() : ((ECKey)jwk).toPrivateKey();
        } catch (JOSEException e) {
            throw new RuntimeException("Failed to get private key", e);
        }

        Map<String, Object> claims = customClaims != null ? new HashMap<>(customClaims) : new HashMap<>();
        return Jwts.builder()
                .setClaims(claims)
                .setSubject(username)
                .setIssuer(issuer)
                .setIssuedAt(new Date())
                .setExpiration(new Date(System.currentTimeMillis() + expiration * 1000))
                .setHeaderParam("kid", jwk.getKeyID())
                .signWith(privateKey)
                .compact();
    }

    public Claims parseToken(String token) {
        JWK jwk = jwkKeyService.getSigningKey();
        PublicKey publicKey;
        try {
            publicKey = jwk instanceof RSAKey ? 
                ((RSAKey)jwk).toPublicKey() : ((ECKey)jwk).toPublicKey();
        } catch (JOSEException e) {
            throw new RuntimeException("Failed to get public key", e);
        }
            
        return Jwts.parserBuilder()
                .setSigningKey(publicKey)
                .build()
                .parseClaimsJws(token)
                .getBody();
    }

    public boolean validateToken(String token) {
        try {
            parseToken(token);
            return true;
        } catch (Exception e) {
            return false;
        }
    }
}
